package com.syher.util.heartbeat;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.ObjectOutputStream;
import java.net.Socket;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author shenyuhang
 * @date 2018/07/03.
 **/
public class HeartbeatSender {

    private static final Logger LOGGER = LoggerFactory.getLogger(HeartbeatSender.class);
    private static final Integer DEFAULT_INTERVAL = 5 * 1000;

    private String host;
    private Integer port;
    private long interval;
    private boolean retry;
    private HeartbeatCallback callback;

    private ScheduledExecutorService executor;
    private Map<String, Boolean> status = new ConcurrentHashMap<>(1);

    public HeartbeatSender(String host, Integer port) {
        this.host = host;
        this.port = port;
        this.interval = DEFAULT_INTERVAL;
        this.retry = false;
    }

    public void setInterval(long interval) {
        this.interval = interval;
    }

    public void setRetry(boolean retry) {
        this.retry = retry;
    }

    public void setCallback(HeartbeatCallback callback) {
        this.callback = callback;
    }

    public void sendHeartbeat() {
        executor = Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder().setNameFormat("heartbeat").build());
        executor.scheduleWithFixedDelay(new HeartbeatRunner(), 0, interval, TimeUnit.MILLISECONDS);
    }

    public void stop() {
        executor.shutdown();
        executor = null;
    }

    private void retry() {
        stop();
        if (retry) {
            sendHeartbeat();
        }
    }

    class HeartbeatRunner implements Runnable {

        AtomicBoolean running;

        @Override
        public void run() {
            Socket socket = null;
            running = new AtomicBoolean(true);
            String hostAndPort = String.join(":", host, port.toString());
            try {
                socket = new Socket(host, port);
                String message = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date()) + "\t keep alive.";
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("send heartbeat message:{}", message);
                }
                ObjectOutputStream oos = new ObjectOutputStream(socket.getOutputStream());
                oos.writeObject(message);
                oos.flush();

                if (status.isEmpty() || status.get(hostAndPort).compareTo(true) != 0) {
                    if (callback != null) {
                        callback.trigger("connected to " + String.join(":", host, port.toString()));
                    }
                }

            } catch (IOException e) {
                if (status.isEmpty() || status.get(hostAndPort).compareTo(false) != 0) {
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.warn("error send heartbeat:{}", e.getMessage());
                    }
                    if (callback != null) {
                        callback.trigger("could not connect to " + String.join(":", host, port.toString()));
                    }
                }

                running.set(false);
                try {
                    Thread.sleep(HeartbeatSender.this.interval);
                    HeartbeatSender.this.retry();
                } catch (InterruptedException e1) {
                    HeartbeatSender.this.stop();
                }
            } finally {
                status.put(hostAndPort, running.get());
                try {
                    if (socket != null) {
                        socket.close();
                    }
                } catch (IOException e) {
                    socket = null;
                    HeartbeatSender.this.stop();
                }
            }
        }
    }

    @FunctionalInterface
    public interface HeartbeatCallback {
        /**
         * 连接状态变更时触发
         */
        void trigger(Object object);
    }
}
